import numpy as np
import cv2
import glob
from math import *
import pandas as pd
import os
import sys
sys.path.append('././')
from kinematic import *
# 标定板 9*6 25mm
chess_board_x_num = 10 - 1
chess_board_y_num = 7 - 1
chess_board_len = 25

# 相机内参矩阵和畸变系数
cam_mtx = np.array( [[670.92263044  ,      0.        ,   358.82135372],
                    [  0.           ,    673.8728403 ,   224.19994634],
                    [  0.           ,      0.        ,     1.        ]], dtype=np.float64)

cam_dist = np.array([-0.19770613, -0.19343703 ,-0.00058339, -0.00063087 , 0.07889203], dtype=np.float64)

#初始化的，数据不用看
cam_end_mtx = np.array([[7.09365529e-01, 7.03646951e-01, -4.10062674e-02, 3.03824839e+01],
[-7.04212502e-01, 7.09988680e-01, 9.09529877e-04, -1.74020390e+01],
[ 2.97539736e-02, 2.82319370e-02, 9.99158475e-01, 6.56932571e+00],
[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]], dtype=np.float64)
#初始化的，数据不用看
chess_base_mtx = np.array([[-9.99352002e-01, -2.72295953e-02, -2.35398822e-02, -7.35039823e+01], 
[2.42275009e-02, -2.52045413e-02, -9.99388693e-01, -5.35886457e+02], 
[2.66196377e-02, -9.99311403e-01, 2.58479139e-02, 4.23558485e+02], 
[0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 1.00000000e+00]], dtype=np.float64)


count = 0

def get_RT_from_chessboard(img_path, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist):
    '''
    #用来从棋盘格图片得到相机外参
    :param img_path: 读取图片路径
    :param chess_board_x_num: 棋盘格x方向格子数
    :param chess_board_y_num: 棋盘格y方向格子数
    :param chess_board_len: 单位棋盘格长度,mm
    :param cam_mtx: 相机内参
    :param cam_dist: 相机畸变参数
    :return: 相机外参
    '''
    global count
    
    # termination criteria
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    # 标定板世界坐标
    objp = np.zeros((chess_board_y_num*chess_board_x_num,3), np.float32)
    for m in range(chess_board_y_num):
        for n in range(chess_board_x_num):
            objp[m*chess_board_x_num + n] = [n*chess_board_len, m*chess_board_len, 0]
    print('obj:\n',objp)
    # Arrays to store object points and image points from all the images.
    objpoints = [] # 3d point in real world space
    imgpoints = [] # 2d points in image plane.

    img = cv2.imread(img_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # Find the chess board corners
    ret, corners = cv2.findChessboardCorners(gray, (chess_board_x_num, chess_board_y_num), None)

    # If found, add object points, image points (after refining them)
    if ret == True:
        objpoints.append(objp)
        corners2 = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)
        imgpoints.append(corners2)
        draw_p=cv2.drawChessboardCorners(img, (chess_board_x_num, chess_board_y_num), corners, ret)
        cv2.imwrite('calibration/handeye_clib/drawChessboardCorners/'+str(count)+'.jpg', draw_p)
        count+=1
        # Draw and display the corners
        # cv2.namedWindow('img', 0)
        # cv2.resizeWindow('img', int(2560*0.6), int(1440*0.6))
        # cv2.drawChessboardCorners(img, (chess_board_x_num, chess_board_y_num), corners2, ret)
        # cv2.imshow('img', img)
        # cv2.waitKey(10)
        # input("请输入任意字符继续：")
        retval, rvec, tvec = cv2.solvePnP(objpoints[0], imgpoints[0], cam_mtx, cam_dist)
        # print(rvec.reshape((1,3)))
        # 反投影误差
        total_error = 0
        imgpoints2, _ = cv2.projectPoints(objpoints[0], rvec, tvec, cam_mtx, cam_dist)
        error = cv2.norm(imgpoints[0], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
        total_error += error
        print("SolvePnP error: ", total_error / len(objpoints))

        RT = np.column_stack(((cv2.Rodrigues(rvec))[0],tvec))
        RT = np.row_stack((RT, np.array([0, 0, 0, 1])))
        #注意要求逆
        # RT=np.linalg.inv(RT)
        # print("RT\n",RT)
        return True,RT
    return False,None


def get_mtx(path, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist):
    img_path =path+ "/*.jpg"
    coord_file_path = path + '/机械臂末端位姿.xlsx' #从记录文件读取机器人六个位姿
    df = pd.read_excel(coord_file_path)
    R_all_end_to_base_1=[]
    T_all_end_to_base_1=[]
    R_all_chess_to_cam_1=[]
    T_all_chess_to_cam_1=[]
    images = glob.glob(img_path)
    for row,fname in enumerate(images):
        print('第'+str(row)+'张图片:')
        ret,RT_chess_to_cam = get_RT_from_chessboard(fname, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist)
        if ret:
            R_all_chess_to_cam_1.append(RT_chess_to_cam[:3,:3])
            T_all_chess_to_cam_1.append(RT_chess_to_cam[:3, 3].reshape((3,1)))
            print('solvePnP成功:')
            print(RT_chess_to_cam)

            joint_angles=[df.iloc[row, 0], df.iloc[row, 1], df.iloc[row, 2], df.iloc[row, 3], df.iloc[row, 4]]
            RT_end_to_base = forward_kinematic(np.deg2rad(joint_angles))
            R_all_end_to_base_1.append(RT_end_to_base[:3, :3])
            T_all_end_to_base_1.append(RT_end_to_base[:3, 3].reshape((3, 1)))
            print('正解成功:')
            print(RT_end_to_base)
            print('')
    
    return R_all_chess_to_cam_1, T_all_chess_to_cam_1,R_all_end_to_base_1, T_all_end_to_base_1


def hangeye_calibration():
    print("开始手眼标定\n")
    path = "calibration/handeye_clib/hand_eye"

    R_all_chess_to_cam_1, T_all_chess_to_cam_1,R_all_end_to_base_1, T_all_end_to_base_1=get_mtx(path, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist)

    good_pictures=len(R_all_chess_to_cam_1)
    if good_pictures<2:
        return good_pictures,None
    R, T = cv2.calibrateHandEye(R_all_end_to_base_1, T_all_end_to_base_1, R_all_chess_to_cam_1, T_all_chess_to_cam_1, cv2.CALIB_HAND_EYE_TSAI)#手眼标定
    RT = np.column_stack((R,T))
    RT = np.row_stack((RT, np.array([0, 0, 0, 1])))#即为cam to end变换矩阵
    print('相机相对于末端的变换矩阵为：')
    print(RT)


    chess_base_mtx = np.zeros((4,4), np.float32)
    print('')
    # 结果验证,原则上来说,每次结果相差较小
    for i in range(len(R_all_chess_to_cam_1)):
        #验证RT_chess_to_base
        RT_end_to_base = np.column_stack((R_all_end_to_base_1[i],T_all_end_to_base_1[i]))
        RT_end_to_base = np.row_stack((RT_end_to_base,np.array([0,0,0,1])))
        
        RT_chess_to_cam = np.column_stack((R_all_chess_to_cam_1[i],T_all_chess_to_cam_1[i]))
        RT_chess_to_cam = np.row_stack((RT_chess_to_cam,np.array([0,0,0,1])))

        RT_cam_to_end = np.column_stack((R,T))
        RT_cam_to_end = np.row_stack((RT_cam_to_end,np.array([0,0,0,1])))
        R_cam_to_end = RT_cam_to_end[:3, :3]
        T_cam_to_end = RT_cam_to_end[:3, 3]
        rotation_vector, _ = cv2.Rodrigues(R_cam_to_end)
        thetaX = rotation_vector[0][0] * 180 / pi
        thetaY = rotation_vector[1][0] * 180 / pi
        thetaZ = rotation_vector[2][0] * 180 / pi
        # print(RT_cam_to_end)

        RT_chess_to_base = RT_end_to_base@RT_cam_to_end@RT_chess_to_cam # 即为固定的棋盘格相对于机器人基坐标系位姿
        # RT_chess_to_base = np.linalg.inv(RT_chess_to_base)
        print('chess_to_base 第',i,'次:')
        print(RT_chess_to_base)

        chess_base_mtx += RT_chess_to_base
        
        #检查棋盘格原点在图片中的像素坐标是否准确
        chess_pose = np.array([0,0,0,1])
        # 转换棋盘坐标到像素坐标
        # world_pose = RT_chess_to_base@chess_pose.T
        # cam_pose = np.linalg.inv(RT_end_to_base@RT_cam_to_end)@world_pose
        # # print(cam_pose)
        cam_pose = RT_chess_to_cam@chess_pose.T
        zero3 = np.zeros((3,1))
        imgTocam_matrix = np.hstack((cam_mtx,zero3)) # 3*4 矩阵
        imgTocam_matrix = np.row_stack((imgTocam_matrix,np.array([0,0,0,1])))
        img_pose = imgTocam_matrix@(cam_pose/cam_pose[2])
        print('img_pose 第',i,'次:')
        print(img_pose)
        print('')
    print('平均chess_base_mtx：\n',chess_base_mtx/len(R_all_chess_to_cam_1))
    return good_pictures,RT


# hangeye_calibration()